import torch.nn as nn
import torch
from ..builder import BACKBONES



@BACKBONES.register_module()
class GTI_VGG16(nn.Module):
    def __init__(self,depth=16, use_bn=False, init_cfg=None, input_size=[448, 448]):
        super(GTI_VGG16, self).__init__()
        self.use_bn = use_bn
        self.net_cfg = [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M']
        self.backbone = self.make_layers()
        

    def make_layers(self):
        layers = []
        in_channels = 3
        out_channels = -1
        for i, out_channels in enumerate(self.net_cfg):
            if out_channels == 'M':
                layers.append(nn.MaxPool2d(kernel_size=[2, 2], stride=[2, 2]))
            else:
                layers.append(nn.Conv2d(in_channels, out_channels, kernel_size = 3, stride=1, padding=1))
                if self.use_bn:
                    layers.append(nn.BatchNorm2d(out_channels))
                layers.append(nn.ReLU())
                in_channels = out_channels
        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.backbone(x)
        outs=[x]
        return tuple(outs)
